package org.west.sky.scripture.ratelimiter;

import org.junit.jupiter.api.Test;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.boot.test.context.SpringBootTest;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.scripting.support.ResourceScriptSource;

import java.time.LocalTime;
import java.util.Collections;

/**
 * @author chz
 * @date 2022/3/21
 * @description: 利用redis zset实现滑动窗口限流, members是当前时间
 */
@SpringBootTest
public class RedisSlideWindowRateLimiter {

    private static String KEY_PREFIX = "limiter_";
    private static String QPS = "3";

    private StringRedisTemplate redisTemplate;

    @Autowired
    public RedisSlideWindowRateLimiter(StringRedisTemplate redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    public boolean acquire(String key) {
        long now = System.currentTimeMillis();
        System.out.print(now + " ");
        key = KEY_PREFIX + key;
        String oldest  = String.valueOf(now - 1000);
        String score = String.valueOf(now);
        String scoreValue = score;
        DefaultRedisScript<Long> redisScript = new DefaultRedisScript<>();
        redisScript.setResultType(Long.class);
        //读取resources目录下的lua文件
        redisScript.setScriptSource(new ResourceScriptSource(new ClassPathResource("redisLua/SlideWindow.lua")));
        Long aLong = redisTemplate.execute(redisScript, Collections.singletonList(key), oldest, score, QPS, scoreValue);
        return aLong != null && aLong == 1;
    }

    @Test
    public void test() throws InterruptedException {
        for (int i = 0; i < 15; i++) {
            Thread.sleep(250);
            if (this.acquire("user")) {
                System.out.println(LocalTime.now() + " 做点什么");
            } else {
                System.out.println(LocalTime.now() + " 被限流了");
            }
        }
    }
}
